/**
 * Copyright (C) 2025 pony working team
 * SPDX-License-Identifier: MulanPSL-2.0
 */

/* includes (local) ----------------------------------------------------------*/
#include <pony/mtd/w25x.h>
#include <hal_spi.h>
/* includes (standard library, system) ---------------------------------------*/
/* includes (other library) --------------------------------------------------*/
/* includes (project) --------------------------------------------------------*/
#include <x_module.h>

#include <k_errno.h>

/* defines -------------------------------------------------------------------*/

#define CHOOSE_BIT_16 16
#define CHOOSE_BIT_8  8

/* typedefs ------------------------------------------------------------------*/

typedef struct w25x {
    hal_mtd_t mtd;
    hal_spi_t *spi;
} w25x_t;

/* macro ---------------------------------------------------------------------*/

#define CHECK_RET_RETURN(ret)                                                                      \
    do {                                                                                           \
        if ((ret) < 0) {                                                                           \
            return ret;                                                                            \
        }                                                                                          \
    } while (0)

/* variables (extern) --------------------------------------------------------*/
/* variables (local) ---------------------------------------------------------*/

static w25x_t s_w25x;

/* variables (global) --------------------------------------------------------*/
/* inlines -------------------------------------------------------------------*/

static inline int __w25x_check_block(size_t startblock, size_t nblocks)
{
    if (startblock >= W25X_NR_BLOCKS || startblock + nblocks > W25X_NR_BLOCKS) {
        return -1;
    }
    return 0;
}

static inline int __w25x_check_sector(size_t idx)
{
    if (idx >= W25X_NR_SECTORS) {
        return -1;
    }
    return 0;
}

/* functions -----------------------------------------------------------------*/

/* spi flash */

static int __spi_flash_send_byte(hal_spi_t *spi, uint8_t send, uint8_t *recv)
{
    uint8_t tmp;

    if (hal_spi_write_read(spi, &send, &tmp, 1, HAL_MAX_TIMEOUT) != 0) {
        return -1;
    }
    if (recv != NULL) {
        *recv = tmp;
    }

    return 0;
}

static int __spi_flash_send_cmd_addr(hal_spi_t *spi, uint8_t cmd, size_t addr)
{
    int ret = 0;

    ret = __spi_flash_send_byte(spi, cmd, NULL);
    CHECK_RET_RETURN(ret);

    /* Send the 24-bit address */
    ret = __spi_flash_send_byte(spi, (addr & 0xFF0000) >> CHOOSE_BIT_16, NULL);
    CHECK_RET_RETURN(ret);
    ret = __spi_flash_send_byte(spi, (addr & 0xFF00) >> CHOOSE_BIT_8, NULL);
    CHECK_RET_RETURN(ret);
    ret = __spi_flash_send_byte(spi, addr & 0xFF, NULL);
    CHECK_RET_RETURN(ret);

    return ret;
}

static int __spi_flash_send_one_cmd(hal_spi_t *spi, uint8_t cmd)
{
    int ret;
    hal_spi_ioctl(spi, HAL_SPI_CMD_ENABLE);
    ret = __spi_flash_send_byte(spi, cmd, NULL);
    hal_spi_ioctl(spi, HAL_SPI_CMD_DISABLE);
    return ret;
}

/* w25x private */

static inline int __w25x_write_enable(w25x_t *w25x)
{
    return __spi_flash_send_one_cmd(w25x->spi, W25X_WRITE_ENABLE);
}

static inline int __w25x_power_down(w25x_t *w25x)
{
    return __spi_flash_send_one_cmd(w25x->spi, W25X_POWER_DOWN);
}

static inline int __w25x_wake_up(w25x_t *w25x)
{
    return __spi_flash_send_one_cmd(w25x->spi, W25X_RELEASE_POWER_DOWN);
}

static void __w25x_wait_write_end(w25x_t *w25x)
{
    uint8_t status = 0;
    hal_spi_t *spi = w25x->spi;

    hal_spi_ioctl(spi, HAL_SPI_CMD_ENABLE);

    (void)__spi_flash_send_byte(spi, W25X_READ_STATUS_REG1, NULL);

    /* Loop as long as the memory is busy with a write cycle */
    do {
        /* Send a dummy byte to generate the clock needed by the FLASH
           and put the value of the status register in status variable */
        if (__spi_flash_send_byte(spi, W25X_DUMMY_BYTE, &status) == -1) {
            break;
        }
    } while ((status & W25X_WIP_FLAG) == W25X_WIP_FLAG); /* Write in progress */

    hal_spi_ioctl(spi, HAL_SPI_CMD_DISABLE);
}

static int __w25x_bulk_erase(w25x_t *w25x)
{
    int ret;

    __w25x_write_enable(w25x);
    ret = __spi_flash_send_one_cmd(w25x->spi, W25X_CHIP_ERASE);
    __w25x_wait_write_end(w25x);

    return ret;
}

static int __w25x_erase_area(w25x_t *w25x, uint8_t cmd, size_t addr)
{
    int ret = 0;
    hal_spi_t *spi = w25x->spi;

    __w25x_write_enable(w25x);
    __w25x_wait_write_end(w25x);

    hal_spi_ioctl(spi, HAL_SPI_CMD_ENABLE);

    ret = __spi_flash_send_cmd_addr(spi, cmd, addr);

    hal_spi_ioctl(spi, HAL_SPI_CMD_DISABLE);

    __w25x_wait_write_end(w25x);

    return ret;
}

static int __w25x_erase_sector(w25x_t *w25x, uint32_t idx)
{
    if (__w25x_check_sector(idx) != 0) {
        return -EINVAL;
    }
    return __w25x_erase_area(w25x, W25X_SECTOR_ERASE, idx * W25X_SECTOR_SIZE);
}

static uint32_t __w25x_get_id(w25x_t *w25x)
{
    uint8_t tmp1 = 0;
    uint8_t tmp2 = 0;
    uint8_t tmp3 = 0;
    hal_spi_t *spi = w25x->spi;

    hal_spi_ioctl(spi, HAL_SPI_CMD_ENABLE);

    if (__spi_flash_send_byte(spi, W25X_JEDEC_ID, NULL) == 0) {
        (void)__spi_flash_send_byte(spi, W25X_DUMMY_BYTE, &tmp1);
        (void)__spi_flash_send_byte(spi, W25X_DUMMY_BYTE, &tmp2);
        (void)__spi_flash_send_byte(spi, W25X_DUMMY_BYTE, &tmp3);
    }

    hal_spi_ioctl(spi, HAL_SPI_CMD_DISABLE);

    return (tmp1 << CHOOSE_BIT_16) | (tmp2 << CHOOSE_BIT_8) | tmp3;
}

static uint32_t __w25x_get_device_id(w25x_t *w25x)
{
    uint8_t tmp = 0;
    hal_spi_t *spi = w25x->spi;

    hal_spi_ioctl(spi, HAL_SPI_CMD_ENABLE);

    if (__spi_flash_send_byte(spi, W25X_MANUFACTURER_ID, NULL) == 0) {
        (void)__spi_flash_send_byte(spi, W25X_DUMMY_BYTE, NULL);
        (void)__spi_flash_send_byte(spi, W25X_DUMMY_BYTE, NULL);
        (void)__spi_flash_send_byte(spi, W25X_DUMMY_BYTE, NULL);
        (void)__spi_flash_send_byte(spi, W25X_DUMMY_BYTE, &tmp);
    }

    hal_spi_ioctl(spi, HAL_SPI_CMD_DISABLE);

    return tmp;
}

static int __w25x_write_page(w25x_t *w25x, size_t pos, size_t nbytes, const uint8_t *buf)
{
    int ret = 0;
    size_t i;
    hal_spi_t *spi = w25x->spi;

    if (nbytes == 0) {
        return 0;
    }

    __w25x_write_enable(w25x);

    hal_spi_ioctl(spi, HAL_SPI_CMD_ENABLE);

    if ((ret = __spi_flash_send_cmd_addr(spi, W25X_PAGE_PROGRAM, pos)) == 0) {
        for (i = 0; i < nbytes; ++i) {
            if ((ret = __spi_flash_send_byte(spi, buf[i], NULL)) != 0) {
                break;
            }
        }
    }

    hal_spi_ioctl(spi, HAL_SPI_CMD_DISABLE);

    __w25x_wait_write_end(w25x);

    return ret;
}

/* w25x api */

static int __w25x_read(hal_mtd_t *mtd, size_t pos, size_t nbytes, uint8_t *buf);
static int __w25x_write(hal_mtd_t *mtd, size_t pos, size_t nbytes, const uint8_t *buf);

static int __w25x_erase_block(hal_mtd_t *mtd, size_t startblock, size_t nblocks)
{
    if (__w25x_check_block(startblock, nblocks) != 0) {
        return -EINVAL;
    }

    size_t i;
    w25x_t *w25x = container_of(mtd, w25x_t, mtd);

    for (i = 0; i < nblocks; ++i) {
        if (__w25x_erase_area(w25x, W25X_BLOCK_ERASE, (startblock + i) * W25X_BLOCK_SIZE) != 0) {
            return -1;
        }
    }

    return 0;
}

static int __w25x_read_block(hal_mtd_t *mtd, size_t startblock, size_t nblocks, uint8_t *buf)
{
    if (__w25x_check_block(startblock, nblocks) != 0) {
        return -EINVAL;
    }

    return __w25x_read(mtd, startblock * W25X_BLOCK_SIZE, nblocks * W25X_BLOCK_SIZE, buf);
}

static int __w25x_write_block(hal_mtd_t *mtd, size_t startblock, size_t nblocks, const uint8_t *buf)
{
    size_t i;

    if (__w25x_check_block(startblock, nblocks) != 0) {
        return -EINVAL;
    }

    w25x_t *w25x = container_of(mtd, w25x_t, mtd);
    size_t pages = nblocks * W25X_BLOCK_SIZE / W25X_PAGE_SIZE;

    for (i = 0; i < pages; i++) {
        size_t step = i * W25X_PAGE_SIZE;
        if (__w25x_write_page(w25x, startblock + step, W25X_PAGE_SIZE, buf + step) != 0) {
            return -1;
        }
    }
    return 0;
}

static int __w25x_erase(hal_mtd_t *mtd, size_t pos, size_t nbytes)
{
    size_t start_sector = pos / W25X_SECTOR_SIZE;
    size_t end_sector = (pos + nbytes - 1) / W25X_SECTOR_SIZE;

    if (__w25x_check_sector(start_sector) != 0 || __w25x_check_sector(end_sector) != 0) {
        return -EINVAL;
    }

    w25x_t *w25x = container_of(mtd, w25x_t, mtd);

    for (size_t i = start_sector; i <= end_sector; ++i) {
        if (__w25x_erase_area(w25x, W25X_SECTOR_ERASE, i * W25X_SECTOR_SIZE) != 0) {
            return -1;
        }
    }

    return 0;
}

static int __w25x_read(hal_mtd_t *mtd, size_t pos, size_t nbytes, uint8_t *buf)
{
    int ret = 0;
    size_t i;
    w25x_t *w25x = container_of(mtd, w25x_t, mtd);
    hal_spi_t *spi = w25x->spi;

    hal_spi_ioctl(spi, HAL_SPI_CMD_ENABLE);

    if ((ret = __spi_flash_send_cmd_addr(spi, W25X_READ_DATA, pos)) == 0) {
        for (i = 0; i < nbytes; ++i) {
            if (__spi_flash_send_byte(spi, W25X_DUMMY_BYTE, buf + i) != 0) {
                ret = -1;
                break;
            }
        }
    }

    hal_spi_ioctl(spi, HAL_SPI_CMD_DISABLE);

    return ret;
}

static int __w25x_write(hal_mtd_t *mtd, size_t pos, size_t nbytes, const uint8_t *buf)
{
    w25x_t *w25x = container_of(mtd, w25x_t, mtd);

    return __w25x_write_page(w25x, pos, nbytes, buf);
}

static int __w25x_ioctl(hal_mtd_t *mtd, int cmd, va_list valist)
{
    int ret = -ENOSYS;
    w25x_t *w25x = container_of(mtd, w25x_t, mtd);

    switch (cmd) {
    case HAL_MTD_CMD_GET_ID: {
        uint32_t *buf = va_arg(valist, uint32_t *);
        if (buf == NULL) {
            return -EINVAL;
        }
        uint32_t id = __w25x_get_id(w25x);
        *buf = id;
        return 0;
    }
    case HAL_MTD_CMD_GET_DEV_ID: {
        uint32_t *buf = va_arg(valist, uint32_t *);
        if (buf == NULL) {
            return -EINVAL;
        }
        uint32_t id = __w25x_get_device_id(w25x);
        *buf = id;
        return 0;
    }
    case HAL_MTD_CMD_LOCK:
        return __w25x_power_down(w25x);
    case HAL_MTD_CMD_UNLOCK:
        return __w25x_wake_up(w25x);
    case HAL_MTD_CMD_CHIP_ERASE:
        return __w25x_bulk_erase(w25x);
    case HAL_MTD_CMD_SECTOR_ERASE: {
        uint32_t idx = va_arg(valist, uint32_t);
        return __w25x_erase_sector(w25x, idx);
    }
    case HAL_MTD_CMD_GET_SECTOR_SIZE: {
        va_arg(valist, size_t);   // eat pos
        size_t *buf = va_arg(valist, size_t *);
        if (buf == NULL) {
            return -EINVAL;
        }
        *buf = W25X_SECTOR_SIZE;
        return 0;
    }
    default:
        break;
    }
    return ret;
}

static const hal_mtd_methods_t w25x_methods = {
    .erase_block = __w25x_erase_block,
    .read_block = __w25x_read_block,
    .write_block = __w25x_write_block,
    .erase = __w25x_erase,
    .read = __w25x_read,
    .write = __w25x_write,
    .ioctl = __w25x_ioctl,
};

static int __w25x_open(hal_dev_t *dev, int oflag, va_list valist)
{
    w25x_t *w25x = (w25x_t *)dev;

    w25x->spi = hal_spi_open(CONFIG_W25X_SPI_BUS);
    if (w25x->spi == NULL) {
        return -1;
    }

    return 0;
}

static hal_dev_ops_t dev_ops = {
    .open = __w25x_open,
};

static int w25x_init(void)
{
    s_w25x.mtd.page_size = W25X_PAGE_SIZE;
    s_w25x.mtd.total_size = W25X_NR_BYTES;
    if (hal_mtd_register(&s_w25x.mtd, "w25x", &w25x_methods, 0)) {
        return -1;
    }
    if (hal_dev_init((hal_dev_t *)&s_w25x, &dev_ops) != 0) {
        return -2;
    }

    return 0;
}

MODULE_INIT(driver, w25x_init);
